// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions {
namespace {

class MapTopNTest : public test::FunctionBaseTest {};

TEST_F(MapTopNTest, basic) {
  auto data = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{1:3, 2:5, 3:1, 4:4, 5:2}",
          "{1:3, 2:5, 3:null, 4:4, 5:2}",
          "{1:null, 2:null, 3:1, 4:4, 5:null}",
          "{1:10, 2:7, 3:11, 5:4}",
          "{1:10, 2:7, 3:0}",
          "{1:null, 2:10}",
          "{}",
          "{1:null, 2:null, 3:null}",
      }),
  });

  auto result = evaluate("map_top_n(c0, 3)", data);

  auto expected = make_map_vector_from_json<int32_t, int64_t>({
      "{1:3, 2:5, 4:4}",
      "{1:3, 2:5, 4:4}",
      "{3:1, 4:4, 5:null}",
      "{1:10, 2:7, 3:11}",
      "{1:10, 2:7, 3:0}",
      "{1:null, 2:10}",
      "{}",
      "{1:null, 2:null, 3:null}",
  });

  assertEqualVectors(expected, result);

  // n = 0. Expect empty maps.
  result = evaluate("map_top_n(c0, 0)", data);

  expected = make_map_vector_from_json<int32_t, int64_t>(
      {"{}", "{}", "{}", "{}", "{}", "{}", "{}", "{}"});

  assertEqualVectors(expected, result);

  // n is negative. Expect an error.
  POLLUX_ASSERT_THROW(
      evaluate("map_top_n(c0, -1)", data),
      "n must be greater than or equal to 0");
}

TEST_F(MapTopNTest, nestedNullFailure) {
  auto data = make_map_vector(
      /*offsets=*/{0},
      /*keyVector=*/make_flat_vector<int32_t>({1, 2, 3}),
      /*value_vector=*/
      make_nullable_array_vector<int32_t>(
          std::vector<std::vector<std::optional<int32_t>>>{
              {std::nullopt}, {2}, {5}}));

  // Nested nulls present inhibit the orderbility of values. Expect an error.
  POLLUX_ASSERT_THROW(
      evaluate(
          "map_top_n(c0, 1)",
          make_row_vector(std::vector<VectorPtr>{data})), // n < map size
      "Ordering nulls is not supported");
  POLLUX_ASSERT_THROW(
      evaluate(
          "map_top_n(c0, 10)",
          make_row_vector(std::vector<VectorPtr>{data})), // n > map size
      "Ordering nulls is not supported");
}

// Test to ensure we use keys to break ties if values are
// equal.
TEST_F(MapTopNTest, equalValues) {
  auto data = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>(
          {"{6:3, 2:5, 3:1, 4:4, 5:2, 1:3}",
           "{1:3, 2:5, 3:null, 4:4, 5:2, 6:5 }",
           "{5:null, 2:null, 3:1, 4:4, 1:null}",
           "{1:null, 5:null, 3:null, 4:null, 2:null}"}),
  });

  auto result = evaluate("map_top_n(c0, 3)", data);

  auto expected = make_map_vector_from_json<int32_t, int64_t>(
      {"{6:3, 2:5, 4:4}",
       "{2:5, 4:4, 6:5}",
       "{5:null, 3:1, 4:4}",
       "{5:null, 3:null, 4:null}"});

  assertEqualVectors(expected, result);

  // Map vector with equal array's as values.
  auto valuesVector = make_array_vector_from_json<int64_t>({
      "[1, 2, 3]",
      "[4, 5, 6]",
      "[-1, -2, -3]",
      "[1, 2, 3]",
  });

  auto keysVector = make_flat_vector<int32_t>({1, 2, 3, 4});

  auto mapvector = make_map_vector({0, 4}, keysVector, valuesVector);

  result = evaluate("map_top_n(c0, 3)", make_row_vector({mapvector}));

  auto expectedValues = make_array_vector_from_json<int64_t>({
      "[1, 2, 3]",
      "[4, 5, 6]",
      "[1, 2, 3]",
  });

  auto expectedKeys = make_flat_vector<int32_t>({1, 2, 4});

  auto expectedResults = make_map_vector({0, 3}, expectedKeys, expectedValues);

  assertEqualVectors(expectedResults, result);

  auto dataWithStrKey =
      make_row_vector({make_map_vector_from_json<std::string, int64_t>(
          {R"({"a":2, "b":3, "c":1})", R"({"a":null, "b":3,
            "c":null})"})});

  auto resultWithStrKey = evaluate("map_top_n(c0, 2)", dataWithStrKey);

  auto expectedWithStrKey = make_map_vector_from_json<std::string, int64_t>(
      {R"({"a":2, "b":3})", R"({"b":3, "c":null})"});
  assertEqualVectors(expectedWithStrKey, resultWithStrKey);
}

TEST_F(MapTopNTest, timestampWithTimeZone) {
  auto testMapTopN =
      [&](const std::vector<int32_t>& keys,
          const std::vector<std::optional<int64_t>>& values,
          const std::vector<int32_t>& expectedKeys,
          const std::vector<std::optional<int64_t>>& expectedValues) {
        const auto map = make_map_vector(
            {0},
            make_flat_vector(keys),
            make_nullable_flat_vector(values, TIMESTAMP_WITH_TIME_ZONE()));
        const auto expectedMap = make_map_vector(
            {0},
            make_flat_vector(expectedKeys),
            make_nullable_flat_vector(expectedValues, TIMESTAMP_WITH_TIME_ZONE()));

        const auto result = evaluate("map_top_n(c0, 3)", make_row_vector({map}));

        assertEqualVectors(expectedMap, result);
      };

  testMapTopN(
      {1, 2, 3, 4, 5},
      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5)},
      {1, 2, 4},
      {
          pack(3, 1),
          pack(5, 2),
          pack(4, 4),
      });
  testMapTopN(
      {1, 2, 3, 4, 5},
      {pack(3, 5), pack(5, 4), std::nullopt, pack(4, 2), pack(2, 1)},
      {1, 2, 4},
      {pack(3, 5), pack(5, 4), pack(4, 2)});
  testMapTopN(
      {1, 2, 3, 4, 5},
      {std::nullopt, std::nullopt, pack(1, 1), pack(4, 4), std::nullopt},
      {3, 4, 5},
      {pack(1, 1), pack(4, 4), std::nullopt});
  testMapTopN(
      {1, 2, 3, 5},
      {pack(10, 1), pack(7, 2), pack(11, 3), pack(4, 4)},
      {1, 2, 3},
      {pack(10, 1), pack(7, 2), pack(11, 3)});
  testMapTopN(
      {1, 2, 3},
      {pack(10, 3), pack(7, 2), pack(0, 1)},
      {1, 2, 3},
      {pack(10, 3), pack(7, 2), pack(0, 1)});
  testMapTopN(
      {1, 2}, {std::nullopt, pack(10, 1)}, {1, 2}, {std::nullopt, pack(10, 1)});
  testMapTopN({}, {}, {}, {});
  testMapTopN(
      {1, 2, 3},
      {std::nullopt, std::nullopt, std::nullopt},
      {1, 2, 3},
      {std::nullopt, std::nullopt, std::nullopt});
  testMapTopN(
      {6, 2, 3, 4, 5, 1},
      {pack(3, 1), pack(5, 2), pack(1, 3), pack(4, 4), pack(2, 5), pack(3, 1)},
      {6, 4, 2},
      {pack(3, 1), pack(4, 4), pack(5, 2)});
  testMapTopN(
      {1, 2, 3, 4, 5, 6},
      {pack(3, 6),
       pack(5, 5),
       std::nullopt,
       pack(4, 3),
       pack(2, 2),
       pack(5, 5)},
      {2, 4, 6},
      {
          pack(5, 5),
          pack(4, 3),
          pack(5, 5),
      });
  testMapTopN(
      {5, 2, 3, 4, 1},
      {std::nullopt, std::nullopt, pack(1, 1), pack(4, 2), std::nullopt},
      {5, 3, 4},
      {std::nullopt, pack(1, 1), pack(4, 2)});
  testMapTopN(
      {1, 5, 3, 4, 2},
      {std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt},
      {5, 3, 4},
      {std::nullopt, std::nullopt, std::nullopt});
}

} // namespace
} // namespace kumo::pollux::functions
